//
// Created by engineai on 2024/07/03.
//

#ifndef ZQ_HUMANOID_SPATIALINERTIA_H
#define ZQ_HUMANOID_SPATIALINERTIA_H

#include <cmath>
#include <iostream>
#include <type_traits>

#include "../Math/orientation_tools.h"
#include "spatial.h"

using namespace ori;
using namespace spatial;

/*!
 * Representation of Rigid Body Inertia as a 6x6 Spatial Inertia Tensor
 */
template <typename T>
class SpatialInertia
{
public:
    EIGEN_MAKE_ALIGNED_OPERATOR_NEW
    /*!
     * Construct spatial inertia from mass, center of mass, and 3x3 rotational
     * inertia
     */
    SpatialInertia(T mass, const Vec3<T> &com, const Mat3<T> &inertia)
    {
        Mat3<T> cSkew = vectorToSkewMat(com);
        _inertia.template topLeftCorner<3, 3>() =
            inertia + mass * cSkew * cSkew.transpose();
        _inertia.template topRightCorner<3, 3>() = mass * cSkew;
        _inertia.template bottomLeftCorner<3, 3>() = mass * cSkew.transpose();
        _inertia.template bottomRightCorner<3, 3>() = mass * Mat3<T>::Identity();
    }

    /*!
     * Construct spatial inertia from 6x6 matrix
     */
    explicit SpatialInertia(const Mat6<T> &inertia) { _inertia = inertia; }

    /*!
     * If no argument is given, zero.
     */
    SpatialInertia() { _inertia = Mat6<T>::Zero(); }

    /*!
     * Construct spatial inertia from mass property vector
     */
    explicit SpatialInertia(const MassProperties<T> &a)
    {
        _inertia(0, 0) = a(4);
        _inertia(0, 1) = a(9);
        _inertia(0, 2) = a(8);
        _inertia(1, 0) = a(9);
        _inertia(1, 1) = a(5);
        _inertia(1, 2) = a(7);
        _inertia(2, 0) = a(8);
        _inertia(2, 1) = a(7);
        _inertia(2, 2) = a(6);
        Mat3<T> cSkew = vectorToSkewMat(Vec3<T>(a(1), a(2), a(3)));
        _inertia.template topRightCorner<3, 3>() = cSkew;
        _inertia.template bottomLeftCorner<3, 3>() = cSkew.transpose();
        _inertia.template bottomRightCorner<3, 3>() = a(0) * Mat3<T>::Identity();
    }

    /*!
     * Construct spatial inertia from pseudo-inertia. This is described in
     * Linear Matrix Inequalities for Physically Consistent Inertial Parameter
     *   Identification: A Statistical Perspective on the Mass Distribution, by
     *   Wensing, Kim, Slotine
     * @param P
     */
    explicit SpatialInertia(const Mat4<T> &P)
    {
        Mat6<T> I;
        T m = P(3, 3);
        Vec3<T> h = P.template topRightCorner<3, 1>();
        Mat3<T> E = P.template topLeftCorner<3, 3>();
        Mat3<T> Ibar = E.trace() * Mat3<T>::Identity() - E;
        I.template topLeftCorner<3, 3>() = Ibar;
        I.template topRightCorner<3, 3>() = vectorToSkewMat(h);
        I.template bottomLeftCorner<3, 3>() = vectorToSkewMat(h).transpose();
        I.template bottomRightCorner<3, 3>() = m * Mat3<T>::Identity();
        _inertia = I;
    }

    /*!
     * Convert spatial inertia to mass property vector
     */
    MassProperties<T> asMassPropertyVector()
    {
        MassProperties<T> a;
        Vec3<T> h = matToSkewVec(_inertia.template topRightCorner<3, 3>());
        a << _inertia(5, 5), h(0), h(1), h(2), _inertia(0, 0), _inertia(1, 1),
            _inertia(2, 2), _inertia(2, 1), _inertia(2, 0), _inertia(1, 0);
        return a;
    }

    /*!
     * Get 6x6 spatial inertia
     */
    const Mat6<T> &getMatrix() const { return _inertia; }

    void setMatrix(const Mat6<T> &mat) { _inertia = mat; }

    void addMatrix(const Mat6<T> &mat) { _inertia += mat; }

    /*!
     * Get mass
     */
    T getMass() { return _inertia(5, 5); }

    /*!
     * Get center of mass location
     */
    Vec3<T> getCOM()
    {
        T m = getMass();
        Mat3<T> mcSkew = _inertia.template topRightCorner<3, 3>();
        Vec3<T> com = matToSkewVec(mcSkew) / m;
        return com;
    }

    /*!
     * Get 3x3 rotational inertia
     */
    Mat3<T> getInertiaTensor()
    {
        T m = getMass();
        Mat3<T> mcSkew = _inertia.template topRightCorner<3, 3>();
        Mat3<T> I_rot = _inertia.template topLeftCorner<3, 3>() -
                        mcSkew * mcSkew.transpose() / m;
        return I_rot;
    }

    /*!
     * Convert to 4x4 pseudo-inertia matrix.  This is described in
     * Linear Matrix Inequalities for Physically Consistent Inertial Parameter
     *   Identification: A Statistical Perspective on the Mass Distribution, by
     *   Wensing, Kim, Slotine
     */
    Mat4<T> getPseudoInertia()
    {
        Vec3<T> h = matToSkewVec(_inertia.template topRightCorner<3, 3>());
        Mat3<T> Ibar = _inertia.template topLeftCorner<3, 3>();
        T m = _inertia(5, 5);
        Mat4<T> P;
        P.template topLeftCorner<3, 3>() =
            0.5 * Ibar.trace() * Mat3<T>::Identity() - Ibar;
        P.template topRightCorner<3, 1>() = h;
        P.template bottomLeftCorner<1, 3>() = h.transpose();
        P(3, 3) = m;
        return P;
    }

    /*!
     * Flip inertia matrix around an axis.  This isn't efficient, but it works!
     */
    SpatialInertia flipAlongAxis(CoordinateAxis axis)
    {
        Mat4<T> P = getPseudoInertia();
        Mat4<T> X = Mat4<T>::Identity();
        if (axis == CoordinateAxis::X)
            X(0, 0) = -1;
        else if (axis == CoordinateAxis::Y)
            X(1, 1) = -1;
        else if (axis == CoordinateAxis::Z)
            X(2, 2) = -1;
        P = X * P * X;
        return SpatialInertia(P);
    }

private:
    Mat6<T> _inertia;
};
#endif // ZQ_HUMANOID_SPATIALINERTIA_H
